No actionable comments were generated in the recent review. 🎉

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

_{Comment @coderabbitai help to get the list of available commands.}

Codecov Report

❌ Patch coverage is 78.94737% with 4 lines in your changes missing coverage. Please review.
✅ Project coverage is 90.73%. Comparing base (b443f77) to head (59c2ac5).
⚠️ Report is 19 commits behind head on main.

@@            Coverage Diff             @@
##             main     #401      +/-   ##
==========================================
- Coverage   90.77%   90.73%   -0.05%     
==========================================
  Files          91       92       +1     
  Lines       10588    11459     +871     
  Branches     1582     1730     +148     
==========================================
+ Hits         9611    10397     +786     
- Misses        686      730      +44     
- Partials      291      332      +41     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

Hi @pic-asso, thanks for the PR! I currently do not have easy access to a multi-gpu setup to verify this behavior. Could you provide some numbers / screenshot /logs or something similar to show that this PR does achieve what it claims?

Don't get me wrong, the reasoning and code all makes sense, its just that I currently have no way of verifying.

Pushed a follow-up commit (59c2ac5) addressing the review feedback plus a real gap I found while benchmarking on real multi-GPU hardware.

What's new in this commit

1. The original fix only pinned the sliced input to the mode-solver's pure_callback, not the callback's own outputs. mode_E_raw, mode_H_raw, and eff_idx come back from the host callback carrying a maximal/single-device GSPMDSharding, which is incompatible with the surrounding x-sharded computation. On 2 GPUs this still crashed inside compute_mode, just one step later than before. Now pinned via the same pin_to_single_device.
2. jax.lax.with_sharding_constraint can't handle that specific sharding representation — it raises Cannot convert GSPMDSharding {maximal device=0} into SdyArray internally via XLA's Shardy partitioner. Switched pin_to_single_device's implementation to jax.device_put, which performs the placement directly and accepts both the already-NamedSharding sliced-input case and the maximal-sharded callback-output case.
3. Shortened the duplicated pinning comment in sources/mode.py / detectors/mode.py per the "overdocumented" nit.
4. Extracted _replicated_mesh_sharding() as a separately-testable pure function to close the patch-coverage gap — the mesh/sharding construction can now be unit-tested directly with a mocked device list, without needing the subprocess-based multi-device harness. Local coverage on sharding.py is now 96% (branch), with the remaining lines being the actual device_put placement call, which genuinely requires distinct physical devices to exercise.

Benchmark (real 2×H100 hardware, non-PDK)

To validate this against real multi-GPU hardware without touching any confidential PDK geometry, I wrote a generic standalone benchmark: a straight strip waveguide (n_core=2.0 / n_clad=1.44, made-up numbers, no tie to any real platform), ~121M cells, with a ModePlaneSource feeding a downstream ModeOverlapDetector — same code path as compute_mode. The script monkey-patches pin_to_single_device to a no-op to simulate pre-#401 behavior on the same checkout, so all three phases below ran from the exact same commit.

Ran 3 phases in one SLURM allocation (1×H100 reference, 2×H100 without the callback-output pinning, 2×H100 with it):

Bit-exact correctness (identical overlap to full float32 precision) between 1-GPU and the fully-patched 2-GPU run — that, plus no longer crashing, is the actual point of this fix.

I'd flag the memory/timing columns above as inconclusive on their own: GPU0 only drops 18.3→17.3GB (~5%), the 2-GPU total (30.5GB) is actually higher than the 1-GPU total (18.3GB), and the split is uneven (17.3 vs 13.2) — at this size the fixed overhead (mesh setup, program binary, small non-sharded state) is a large enough fraction of the total that it can obscure a real halving one way or the other. Wall-clock is flat-to-slightly-worse too (~168s vs ~171s).

To get a cleaner read, I re-ran the same scene ~1.6× larger (NX 2500→4000, ~194M cells instead of ~121M):

Here the 2-GPU split is genuinely even (17.51 vs 17.51 GB, not 17.3 vs 13.2) and both devices sit below the 1-GPU figure — a much more convincing sharded-not-replicated signature, even though the wall-clock still doesn't improve (same story: fixed per-step sync + the serial callback dominate). I don't want to overclaim a big memory win from either benchmark, but the larger one is a materially cleaner demonstration than the original table, so pairing them here for completeness.

Edit: removed a closing claim here about PhasorDetector needing separate handling — a follow-up reproduction showed that was wrong (a scene with a PhasorDetector runs cleanly under this fix; the crash I'd seen earlier was this same callback-output issue, not PhasorDetector). No action needed there.

Final wrap-up: while chasing the above I also found that gradient-mode multi-GPU (both method="reversible" and "checkpointed") doesn't shard its backward-pass state at all — replicates the full buffer on every device instead. That's unrelated to this PR's scope (mode-solver pure_callback sharding) and doesn't touch PhasorDetector either, so I filed it separately as #417 rather than folding it in here. From my side this PR is done — the fix, tests, and benchmarks above are the full extent of what I have.

Hi @pic-asso, the pre-commit checks are currently failing an the PR needs some re-work to remove the AI bloat. From a functional point of view this looks good to me